import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from skimage import io
import torchvision.transforms as transforms
import torchvision.models as models
from PIL import Image
import time
import MyNet

# 定义数据预处理的转换
mean = [0.5, 0.5, 0.5]
# std = [0.25, 0.25, 0.25]
std = [0.5, 0.5, 0.5]
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean, std),
])
num_classes = 75
label_Butterfly = {
    0:'ADONIS',
    1:'AFRICAN GIANT SWALLOWTAIL',
    2:'AMERICAN SNOOT',
    3:'AN 88',
    4:'APPOLLO',
    5:'ATALA',
    6:'BANDED ORANGE HELICONIAN',
    7:'BANDED PEACOCK',
    8:'BECKERS WHITE',
    9:'BLACK HAIRSTREAK',
    10:'BLUE MORPHO',
    11:'BLUE SPOTTED CROW',
    12:'BROWN SIPROETA',
    13:'CABBAGE WHITE',
    14:'CAIRNS BIRDWING',
    15:'CHECQUERED SKIPPER',
    16:'CHESTNUT',
    17:'CLEOPATRA',
    18:'CLODIUS PARNASSIAN',
    19:'CLOUDED SULPHUR',
    20:'COMMON BANDED AWL',
    21:'COMMON WOOD-NYMPH',
    22:'COPPER TAIL',
    23:'CRECENT',
    24:'CRIMSON PATCH',
    25:'DANAID EGGFLY',
    26:'EASTERN COMA',
    27:'EASTERN DAPPLE WHITE',
    28:'EASTERN PINE ELFIN',
    29:'ELBOWED PIERROT',
    30:'GOLD BANDED',
    31:'GREAT EGGFLY',
    32:'GREAT JAY',
    33:'GREEN CELLED CATTLEHEART',
    34:'GREY HAIRSTREAK',
    35:'INDRA SWALLOW',
    36:'IPHICLUS SISTER',
    37:'JULIA',
    38:'LARGE MARBLE',
    39:'MALACHITE',
    40:'MANGROVE SKIPPER',
    41:'MESTRA',
    42:'METALMARK',
    43:'MILBERTS TORTOISESHELL',
    44:'MONARCH',
    45:'MOURNING CLOAK',
    46:'ORANGE OAKLEAF',
    47:'ORANGE TIP',
    48:'ORCHARD SWALLOW',
    49:'PAINTED LADY',
    50:'PAPER KITE',
    51:'PEACOCK',
    52:'PINE WHITE',
    53:'PIPEVINE SWALLOW',
    54:'POPINJAY',
    55:'PURPLE HAIRSTREAK',
    56:'PURPLISH COPPER',
    57:'QUESTION MARK',
    58:'RED ADMIRAL',
    59:'RED CRACKER',
    60:'RED POSTMAN',
    61:'RED SPOTTED PURPLE',
    62:'SCARCE SWALLOW',
    63:'SILVER SPOT SKIPPER',
    64:'SLEEPY ORANGE',
    65:'SOOTYWING',
    66:'SOUTHERN DOGFACE',
    67:'STRAITED QUEEN',
    68:'TROPICAL LEAFWING',
    69:'TWO BARRED FLASHER',
    70:'ULYSES',
    71:'VICEROY',
    72:'WOOD SATYR',
    73:'YELLOW SWALLOW TAIL',
    74:'ZEBRA LONG WING'
}
device1 = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
print(device1)
train_pth = r'archive\Testing_set.csv'
train_root_dir = r'archive\test'
class butterfly_dataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        self.annotations = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform
        # self.label_encoder = LabelEncoder()
        # self.annotations['label_encoded'] = self.label_encoder.fit_transform(self.annotations['label'])

    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, index):
        img_path = os.path.join(self.root_dir, self.annotations.iloc[index, 0])
        image = io.imread(img_path)
        image = Image.fromarray(image)
        # y_label = torch.tensor(int(self.annotations.iloc[index, 2]))
        if self.transform:
            image = self.transform(image)

        return (image)

parameter = '_'+str(30)+'_'+str(32)+'_'+str(0.0005)+'_'+str(0.9)
def get_alexnet():
    model_alexnet = models.alexnet(weights=None)
    model_alexnet.classifier[6] = nn.Linear(model_alexnet.classifier[6].in_features, num_classes)
    model_alexnet.load_state_dict(torch.load('models/butterfly_AlexNet'+parameter+'.pth', map_location=device1))
    model_alexnet.eval()
    return model_alexnet
def get_VGG():
    model = models.vgg19(weights=None)
    model.classifier[6] = nn.Linear(model.classifier[6].in_features, num_classes)
    model.load_state_dict(torch.load('models/butterfly_VGG19'+parameter+'.pth', map_location=device1))
    model.eval()
    return model
def get_ResNet50():
    model = models.resnet50(weights=None)
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    model.load_state_dict(torch.load('models/butterfly_ResNet50'+parameter+'.pth', map_location=device1))
    model.eval()
    return model
def get_DenseNet121():
    model = models.densenet121(weights=None)
    num_features = model.classifier.in_features
    model.classifier = nn.Linear(num_features, num_classes)
    model.load_state_dict(torch.load('models/butterfly_DenseNet121'+parameter+'.pth', map_location=device1))
    model.eval()
    return model

def get_SqueezeNet():
    model = models.squeezenet1_1(weights=None)
    # 获取模型最后一个层的输入特征数
    num_features = model.classifier[1].in_channels
    # 替换最后一个全连接层
    model.classifier[1] = nn.Conv2d(num_features, num_classes, kernel_size=(1, 1), stride=(1, 1))
    # 将模型的最后一层设置为分类层
    model.num_classes = num_classes
    model.load_state_dict(torch.load('models/butterfly_SqueezeNet'+parameter+'.pth', map_location=device1))
    model.eval()
    return model

def get_MyNet():
    model = MyNet.MyModel1()
    num_features = model.classifier.in_features
    model.classifier = nn.Linear(num_features, num_classes)
    model.load_state_dict(torch.load('models/butterfly_DenseNet121'+parameter+'.pth', map_location=device1))
    model.eval()
    return model

def predict_image(model, dataloader):
    predictions = []
    num_batches = len(dataloader)
    start_time = time.time()
    with torch.no_grad():
        for i, images in enumerate(dataloader, 1):
            # 假设 images 是未分类的蝴蝶图片数据
            outputs = model(images)  # 使用模型进行预测
            _, predicted = torch.max(outputs, 1)  # 获取预测结果（假设是分类的索引）
            predictions.extend(predicted.tolist())
            progress = i / num_batches * 100
            print(f'Progress: {progress:.2f}%')
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"运行时间为：{execution_time} 秒")
    return predictions

if __name__ == '__main__':
    model = get_alexnet()
    # 100张未分类的蝴蝶图片
    test_dataset = butterfly_dataset(train_pth, train_root_dir, transform=transform)
    batch_size = 32
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # 使用模型进行预测
    predictions = predict_image(model, test_dataloader)
    predicted_categories = [label_Butterfly[label] for label in predictions]

    test_labels = pd.read_csv(train_pth)
    # 将预测结果添加到测试集标签数据的第二列，并命名为'label'
    test_labels.insert(1, 'label', predicted_categories)
    # 检查文件夹是否存在，如果不存在则创建
    folder_path = 'evaluate'
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    # 将修改后的测试集标签数据保存为原始的CSV文件中
    test_labels.to_csv(r'evaluate\AlexNet_testing_set5.csv', index=False)


